#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cmath>

#include "NNetContainer.h"



using namespace std;

int perceptronType = 2; //QP
int mode = 1;
string trainFile;
string testFile;
ofstream * testOut;
string testOutFile;
int nLayer;
vector<int> hiddenNum;
double learnRate;
int labelNum;
int ninput;
int trainTms;

void PrintHelp()
{
    cout << "nnet [mode] [nLayer] [l1_num] [l2_num] ... [ln_num] [nDim] [LearnRate] [TrainFile] [TestFile] [OutFile] [TrainTms] [nLabel(mode 1)]" << endl;
    cout << "[mode]           1:Classifier; 2:Regression" << endl;
    cout << "[nLayer]         The number of hidden layer" << endl;
    cout << "[li_num]         The number of perceptron in layer[i]" << endl;
    cout << "[nDim]           The number of input dimension" << endl;
    cout << "[LearnRate]      Learning rate " << endl;
    cout << "[TrainFile]      Training data file name" << endl;
    cout << "[TestFile]       Testing data file name" << endl;
    cout << "[OutFile]        The result file" << endl;
    cout << "[TrainTms]       The number of training rounds" << endl;
    cout << "[nLabel]         The number of label (mode 1 only)" << endl;
}

void ParseArg(char argc, char ** argv)
{
    int t = 1;
    mode = atoi(argv[t++]);
    nLayer = atoi(argv[t++]);
    for (int i = 0; i < nLayer; ++i)
    {
        int n = atoi(argv[t++]);
        hiddenNum.push_back(n);
    }
    ninput = atoi(argv[t++]);
    learnRate = atof(argv[t++]);
    trainFile = argv[t++];
    testFile = argv[t++];
    testOutFile = argv[t++];
    testOut = new ofstream(testOutFile.c_str());
    trainTms = atoi(argv[t++]);
    if (mode == 1)
    {
        labelNum = atoi(argv[t++]);
    }
}

void HandleClassifier()
{
    NNetContainer net(ninput,nLayer,labelNum,hiddenNum,learnRate);
    net.LoadData(trainFile,ninput,labelNum);
    net.Train(trainTms);
    net.Test(testFile,(*testOut));
}

void HandleRegression()
{
    NNetContainer net(ninput,nLayer,1,hiddenNum,learnRate);
    net.LoadData(trainFile,ninput);
    net.Train(trainTms);
    net.Test(testFile,(*testOut));
}

int main(int argc, char ** argv)
{
    if (argc == 1)
    {
        PrintHelp();
        return 0;
    }
    ParseArg(argc, argv);
    if (mode == 1)
    {
        HandleClassifier();
    }
    else 
    {
        HandleRegression();
    }
    testOut->close();
    delete testOut;
}